package grpc

import (
	"context"
	"fmt"

	"gitee.com/zawei/yggdrasil/pkg/errors"
	"gitee.com/zawei/yggdrasil/pkg/trace"
	"github.com/opentracing/opentracing-go"
	"github.com/opentracing/opentracing-go/ext"
	"github.com/opentracing/opentracing-go/log"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

func extractRpcServerOpt(ctx context.Context) opentracing.StartSpanOption {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		md = metadata.New(nil)
	}
	sc, err := opentracing.GlobalTracer().Extract(opentracing.HTTPHeaders, trace.TextArrayMapCarrier(md))
	if err != nil {
		return trace.NilStartSpanOption{}
	}
	return ext.RPCServerOption(sc)
}

type openTracingServerStream struct {
	grpc.ServerStream
	ctx context.Context
}

func (ss *openTracingServerStream) Context() context.Context {
	return ss.ctx
}

func traceUnaryServerInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
	span, ctx := opentracing.StartSpanFromContext(
		ctx,
		info.FullMethod,
		extractRpcServerOpt(ctx),
	)
	ext.SpanKind.Set(span, "grpcServer.unary")
	ext.Component.Set(span, "grpc")
	defer span.Finish()

	resp, err := handler(ctx, req)

	if err != nil {
		gs := errors.FromProto(errors.FromError(err).Proto())
		span.SetTag("code", gs.HttpCode())
		ext.Error.Set(span, true)
		span.LogFields(log.String("event", "error"), log.String("message", fmt.Sprintf("%+v", err)))
	}
	return resp, err
}

func traceStreamServerInterceptor(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	span, ctx := opentracing.StartSpanFromContext(
		ss.Context(),
		info.FullMethod,
		extractRpcServerOpt(ss.Context()),
	)
	ext.SpanKind.Set(span, "grpcServer.stream")
	ext.Component.Set(span, "grpc")
	defer span.Finish()

	return handler(srv, &openTracingServerStream{
		ServerStream: ss,
		ctx:          ctx,
	})
}
